from math import *

action = ('F', 'L', 'R')
state = (0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15)

policy = [{'F':0,'L':1,'R':0,},{'F':1,'L':0,'R':0},{'F':0,'L':0,'R':1},{'F':0,'L':0,'R':1},{'F':1,'L':0,'R':0},{'F':0,'L':1,'R':0},{'F':1,'L':0,'R':0},{'F':0,'L':1,'R':0},{'F':1,'L':0,'R':0},{'F':0,'L':0,'R':1},{'F':0,'L':1,'R':0},{'F':0,'L':1,'R':0},{'F':0,'L':0,'R':1},{'F':1,'L':0,'R':0},{'F':0,'L':1,'R':0},{'F':1,'L':0,'R':0}]


epsi = 0.0000001      

def policy_evaluation(V_pi_new) :
    delta = 1    
    while delta>epsi :
        V_pi = V_pi_new
        P_pi = [[0 for i in range(len(state))] for i in range(len(state))]
        M_pi = [[0] for i in range(len(state))]
        for i in range(len(policy)) :  
            for j in range(len(state)) :
                P_pi[i][j] = sum(policy[i][action[a]] * P(state[i],state[j],action[a]) for a in range(len(action)))
            M_pi[i][0] = sum(policy[i][action[a]] * sum( P(state[i], s_j, action[a]) * R(state[i], s_j) \
            for s_j in range(len(state)) ) for a in range(len(action)))
        V_pi_new = add(M_pi, mult(skalMult(0.8,P_pi),V_pi))
        delta = maxSub(V_pi_new,V_pi)
        #print "Delta = ", delta
    #print V_pi_new
    return V_pi_new
    
def policy_improvement(V_pi) :
    #new_policy = [0 for i in range(len(policy))]
    global policy
    new_policy = [{'F':0,'L':0,'R':0} for i in range(len(policy))]
    policy_stable = 1
    for i in range(len(policy)):
        policy_inv = dict([[v, k] for [k, v] in policy[i].items()])
        b = policy_inv[max(policy_inv.keys())]
        maxActionList = []
        for a in range(len(action)):
            maxActionList.append( sum( P(state[i], s_j, action[a]) * (R(state[i], s_j) + 0.8*V_pi[s_j][0]) \
            for s_j in range(len(state))))
        pi = policy[i].keys()[maxActionList.index(max(maxActionList))]
        if pi != b : policy_stable = 0
        new_policy[i][pi]=1
        #new_policy[i] = pi
    #print "stable: ",policy_stable
    policy = new_policy   
    if policy_stable :
        return 1
    else :
        return 0
    
def policy_iteration() :    
    V = [[0] for i in range(len(state))]
    stable = 0
    while stable == 0 :
        V_new = policy_evaluation(V)
        V = V_new
        stable = policy_improvement(V)
    return policy
       
def P(s_i, s_j, a) :
    if s_i==0 and s_j==0 and a=='F' : return 1
    elif s_i==0 and s_j==3 and a=='L' : return 1
    elif s_i==0 and s_j==1 and a=='R' : return 1
#------------------------------------------------
    elif s_i==1 and s_j==5 and a=='F' : return 1
    elif s_i==1 and s_j==0 and a=='L' : return 1
    elif s_i==1 and s_j==2 and a=='R' : return 1
#------------------------------------------------
    elif s_i==2 and s_j==2 and a=='F' : return 1
    elif s_i==2 and s_j==1 and a=='L' : return 1
    elif s_i==2 and s_j==3 and a=='R' : return 1
#------------------------------------------------
    elif s_i==3 and s_j==3 and a=='F' : return 1
    elif s_i==3 and s_j==2 and a=='L' : return 1
    elif s_i==3 and s_j==0 and a=='R' : return 1
#------------------------------------------------
    elif s_i==4 and s_j==4 and a=='F' : return 1
    elif s_i==4 and s_j==7 and a=='L' : return 1
    elif s_i==4 and s_j==5 and a=='R' : return 1
#------------------------------------------------
    elif s_i==5 and s_j==9 and a=='F' : return 1
    elif s_i==5 and s_j==4 and a=='L' : return 1
    elif s_i==5 and s_j==6 and a=='R' : return 1
#------------------------------------------------
    elif s_i==6 and s_j==6 and a=='F' : return 1
    elif s_i==6 and s_j==5 and a=='L' : return 1
    elif s_i==6 and s_j==7 and a=='R' : return 1
#------------------------------------------------
    elif s_i==7 and s_j==3 and a=='F' : return 1
    elif s_i==7 and s_j==6 and a=='L' : return 1
    elif s_i==7 and s_j==4 and a=='R' : return 1
#------------------------------------------------
    elif s_i==8 and s_j==8 and a=='F' : return 1
    elif s_i==8 and s_j==11 and a=='L' : return 1
    elif s_i==8 and s_j==9 and a=='R' : return 1
#------------------------------------------------
    elif s_i==9 and s_j==13 and a=='F' : return 1
    elif s_i==9 and s_j==8 and a=='L' : return 1
    elif s_i==9 and s_j==10 and a=='R' : return 1
#------------------------------------------------
    elif s_i==10 and s_j==10 and a=='F' : return 1
    elif s_i==10 and s_j==9 and a=='L' : return 1
    elif s_i==10 and s_j==11 and a=='R' : return 1
#------------------------------------------------
    elif s_i==11 and s_j==7 and a=='F' : return 1
    elif s_i==11 and s_j==10 and a=='L' : return 1
    elif s_i==11 and s_j==8 and a=='R' : return 1
#------------------------------------------------
    elif s_i==12 and s_j==12 and a=='F' : return 1
    elif s_i==12 and s_j==15 and a=='L' : return 1
    elif s_i==12 and s_j==13 and a=='R' : return 1
#------------------------------------------------
    elif s_i==13 and s_j==13 and a=='F' : return 1
    elif s_i==13 and s_j==12 and a=='L' : return 1
    elif s_i==13 and s_j==14 and a=='R' : return 1
#------------------------------------------------
    elif s_i==14 and s_j==14 and a=='F' : return 1
    elif s_i==14 and s_j==13 and a=='L' : return 1
    elif s_i==14 and s_j==15 and a=='R' : return 1
#------------------------------------------------
    elif s_i==15 and s_j==11 and a=='F' : return 1
    elif s_i==15 and s_j==14 and a=='L' : return 1
    elif s_i==15 and s_j==12 and a=='R' : return 1
    else : return 0
    
def R(s_i, s_j) :
    if s_i == s_j : return -1
    elif s_j == 15 : return 1
    else : return 0
    
def zero(m,n):
    new_matrix = [[0 for row in range(n)] for col in range(m)]
    return new_matrix
    
def mult(matrix1,matrix2):
    if len(matrix1[0]) != len(matrix2):
        print 'Matrices must be m*n and n*p to multiply!'
    else:
        new_matrix = zero(len(matrix1),len(matrix2[0]))
        for i in range(len(matrix1)):
            for j in range(len(matrix2[0])):
                for k in range(len(matrix2)):
                    new_matrix[i][j] += matrix1[i][k]*matrix2[k][j]
        return new_matrix
        
def add(matrix1,matrix2):
    if len(matrix1) != len(matrix2) or len(matrix1[0]) != len(matrix2[0]):
        print 'Matrix must be the same length'
    else:
        new_matrix = zero(len(matrix1),len(matrix2[0]))
        for i in range(len(matrix1)):
            for j in range(len(matrix1[0])):
                new_matrix[i][j] = matrix1[i][j] + matrix2[i][j]
        return new_matrix
        
def skalMult(skalar,matrix):
    new_matrix = zero(len(matrix),len(matrix[0]))
    for i in range(len(matrix)):
        for j in range(len(matrix[0])):
            new_matrix[i][j]=matrix[i][j]*skalar
    return new_matrix 
        
def maxSub(matrix1,matrix2):
    max = 0
    if len(matrix1) != len(matrix2) or len(matrix1[0]) != len(matrix2[0]):
        print 'Matrix must be the same length'
    else:
        for i in range(len(matrix1)):
            for j in range(len(matrix1[0])):
                if max < fabs(matrix1[i][j] - matrix2[i][j]) :
                    max = fabs(matrix1[i][j] - matrix2[i][j])
        return max

def print_dict(dicti) :
    for i in range(len(dicti)):
        print "State ", i, " = ", dicti[i]


print_dict(policy_iteration())
